/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: TbeVectorReduce类的实现文件
 * Author: qingjicheng
 * Create: 2020-12-1
 */

#ifndef TBE_VECTOR_REDUCE_H
#define TBE_VECTOR_REDUCE_H

#include <string>
#include <unordered_map>
#include <nlohmann/json.hpp>
#include <hccl/hccl_types.h>
#include <mutex>
#include <map>
#include <list>
#include <hccl/base.h>
#include "op_tiling.h"
#include "private_types.h"
#ifndef HCCD
#include "runtime/rt.h"
#include "graph/operator.h"
#include "graph/compute_graph.h"
#include "graph/node.h"
#endif

namespace TbeReduce {
constexpr u32 TILING_DATA_MAX_SZIE = 32;

struct TilingInputInfo {
    HcclDataType dataType;
    HcclReduceOp redOp;
    u64 count;
    bool operator == (const TilingInputInfo& other) const
    {
        return dataType == other.dataType && redOp == other.redOp && count && other.count;
    }
};

struct TilingHashFuc {
    std::size_t operator()(const TilingInputInfo& key) const
    {
        return std::hash<u64>()(key.count) ^ std::hash<s32>()(key.dataType) ^ std::hash<s32>()(key.redOp);
    }
};

const std::map<std::string, std::string> OPINFO_INDEX_SOC_STRING_MAP = {
    {"Ascend910B1", "_910B1"},
    {"Ascend910B2", "_910B2"},
    {"Ascend910B2C", "_910B2"},
    {"Ascend910B3", "_910B3"},
    {"Ascend910B4", "_910B4"},
    {"Ascend910B4-1", "_910B4-1"},
    {"Ascend910_9391", "_910B1"},
    {"Ascend910_9381", "_910B2"},
    {"Ascend910_9392", "_910B1"},   // Ascend910_9392、Ascend910_9382为预留类型，当前版本暂不支持，待跟随后续版本节奏交付
    {"Ascend910_9382", "_910B2"},
    {"Ascend910_9372", "_910B3"},
    {"Ascend910_9362", "_910B4-1"}
};

class TbeVectorReduce {
public:
    explicit TbeVectorReduce();
    virtual ~TbeVectorReduce();
    HcclResult Init();
    HcclResult Run(const void *src1, const void *src2, u64 count, const HcclDataType dataType,
                               HcclReduceOp redOp, void *stream, const void *dst);
    HcclResult SetGlobalWorkSpace(std::vector<void *> &globalWorkSpaceAddr);
    HcclResult GetOpBinaryPath(std::string &binaryPath);
    HcclResult GetOpProperty(nlohmann::json &opInfo, const std::string propName, std::string &propValue);
    HcclResult GetOpProperty(nlohmann::json &opInfo, const std::string propName, u32 &propValue);
    HcclResult LoadOpBinary();
    HcclResult GetstubFunc(std::string &opName, s8 *&stubFunc, std::string &kernelName);
    HcclResult GetTilingDataDevMem(OpRunInfo &runInfo, void **tilingDataDevPtr,
        const u32 tilingDataMaxSize);

    HcclResult JudgeIsTail(const void *src1, const void *src2, const void *dst, u64 count, const HcclDataType dataType,
        u64 &headCount, u64 &tailCount, void *&tailSrc1, void *&tailSrc2, void *&tailDst);
    HcclResult PrepareVectorReduce(const void *src1, const void *src2, u64 count, const HcclDataType dataType,
        HcclReduceOp redOp, const void *dst, void *stream, void *&addrListDevMemPtr, void *&argsHandle,
        void *&stubFuncAddr, uint32_t &blockDim);
    HcclResult PrepareTailVectorReduce(void *&devMem1, void *&devMem2, u64 tailCount, const HcclDataType dataType,
        HcclReduceOp redOp, void *stream, void *&addrListDevMemPtr, void *&argsHandle, void *&stubFuncAddr,
        uint32_t &blockDim);
    HcclResult JudgeIsSupported(const HcclDataType dataType, const HcclReduceOp redOp);
#ifndef HCCD
    HcclResult AddDataNodeForOriginNode(ge::ComputeGraphPtr &graph, ge::NodePtr &origin_node, size_t input_size);
    HcclResult ConvertToOpRunInfo(gert::TilingContext *context, OpRunInfo &runInfo);
#endif
    bool SupportRt2()
{
    return deviceType_ == DevType::DEV_TYPE_910B || deviceType_ == DevType::DEV_TYPE_910_93;
}
protected:
    BinHandle binHandle_;
    char *tilingDataHostPtr_;
    char *binaryDataPtr_;
    std::mutex stubFuncMutex_;
    std::mutex tilingDataMutex_;
    std::map<std::string, s8 *> opNameStubFuncsMap_;
    bool isInit_;
    DevType deviceType_;

private:
    void InitOpInfoMap310P3();
    void InitOpInfoMap910A();
    void InitOpInfoMap910B();
    HcclResult LoadOpBinaryWithParams(const void *&binaryData, uint64_t &binaryLength);
    HcclResult GenerateStubFunc(s8 *&stubFunc, std::string &kernelName);
    HcclResult VectorReduce(const void *src1, const void *src2, u64 count, const HcclDataType dataType,
                            HcclReduceOp redOp, void *stream, const void *dst);
    HcclResult VectorReduceLoop(const void *src1, const void *src2, u64 count, const HcclDataType dataType,
                            HcclReduceOp redOp, void *stream, const void *dst);
    HcclResult TailVectorReduce(const void *tailSrc1, const void *tailSrc2, u64 tailCount, const HcclDataType dataType,
                                HcclReduceOp redOp, void *stream, void *tailDst);
    HcclResult ExecuteKernelLaunch(const void *src1, const void *src2, const void *dst,
                                   void *stream, const void *stubFunc, OpRunInfo &runInfo, void *tilingDataDevPtr);
    HcclResult GetTilingRunInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo, u64 count,
                                HcclDataType dataType, OpRunInfo &runInfo, HcclReduceOp redOp = HCCL_REDUCE_RESERVED);
    HcclResult RegisterMetadata(nlohmann::json &opInfo);
    bool IsSupportOperatorOverflowDevType910A(HcclDataType dataType, HcclReduceOp redOp) const;
    bool IsSupportOperatorDevType910A(HcclDataType dataType, HcclReduceOp redOp) const;
    HcclResult GetKernelName(nlohmann::json &opInfo, HcclDataType dataType, std::string &kernelName,
                             OpRunInfo& runInfo, HcclReduceOp redOp = HCCL_REDUCE_RESERVED);
    HcclResult GetOpInfo(HcclDataType dataType, HcclReduceOp redOp, nlohmann::json &opDescInfo,
                         nlohmann::json &opTilingInfo);
    HcclResult GetOpInfoIndex(std::string &opInfoIndex, HcclDataType dataType);
    HcclResult GetOpInfoIndex(std::string &opInfoIndex, HcclReduceOp redOp);
    HcclResult GetOpInfoIndex(std::string &opInfoIndex);
    HcclResult GetOpInfoIndex(std::string &opInfoIndex, HcclDataType dataType, HcclReduceOp redOp);
    HcclResult GetTilingDataInfo(const u64 count, const HcclDataType dataType,
        const HcclReduceOp redOp, void*& tilingDataDevPtr, nlohmann::json& opDescInfo, OpRunInfo& runInfo);
    HcclResult GetKernelFunction(const HcclDataType dataType, const HcclReduceOp redOp,
        nlohmann::json& opDescInfo, OpRunInfo& runInfo, s8*& stubFunc);
    HcclResult DevMemMalloc(void* stream, void*& devMem1, void*& devMem2);
    HcclResult ConcatMem(const void *src1, const void *src2, const void *dst, const void *tilingDataDevPtr,
        std::vector<const void *> &deviceAddrs);
#ifndef HCCD
    HcclResult GetTilingRunInfo(nlohmann::json &opDescInfo, u64 count,
                            HcclDataType dataType, OpRunInfo &runInfo, HcclReduceOp redOp = HCCL_REDUCE_RESERVED);
    HcclResult GetOpInfo(HcclDataType dataType, HcclReduceOp redOp, nlohmann::json &opDescInfo);
#endif
    HcclResult GetCoreNum();
    std::map<std::string, nlohmann::json> opInfoMap_;
    std::map<int32_t, void *> devMemMap_; // streamId和device内存的map
    std::mutex devMemMutex_;
    std::unordered_map<TilingInputInfo, void *, TilingHashFuc> tilingDataMap_;
    std::vector<void *> globalWorkSpaceAddr_;
};
} // namespace TbeReduce

#endif
